
package com.sskj.ws.config;

import org.apache.commons.lang.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptorAdapter;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.web.socket.config.annotation.AbstractWebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import com.sskj.security.auth.util.JwtTokenUtil;

@Configuration
@EnableWebSocketMessageBroker
public class MyWebSocketConfig
		extends AbstractWebSocketMessageBrokerConfigurer {

	@Autowired
	private UserDetailsService userDetailsService;

	@Autowired
	private JwtTokenUtil jwtTokenUtil;

	@Value("${jwt.header}")
	private String tokenHeader;

	@Override
	public void configureClientInboundChannel(
			ChannelRegistration registration) {

		registration.setInterceptors(new ChannelInterceptorAdapter() {

			@Override
			public Message<?> preSend(Message<?> message,
					MessageChannel channel) {

				StompHeaderAccessor accessor = MessageHeaderAccessor
						.getAccessor(message, StompHeaderAccessor.class);
				if (StompCommand.CONNECT.equals(accessor.getCommand())) {
					String jwtToken = accessor
							.getFirstNativeHeader("Authorization");

					if (StringUtils.isNotEmpty(jwtToken)) {

						String username = jwtTokenUtil
								.getUsernameFromToken(jwtToken);

						UserDetails userDetails = userDetailsService
								.loadUserByUsername(username);

						if (username != null && SecurityContextHolder
								.getContext().getAuthentication() == null) {

							if (jwtTokenUtil.validateToken(jwtToken,
									userDetails)) {
								UsernamePasswordAuthenticationToken authentication = new UsernamePasswordAuthenticationToken(
										userDetails, null,
										userDetails.getAuthorities());

								SecurityContextHolder.getContext()
										.setAuthentication(authentication);

								accessor.setUser(authentication);
							}
						}
					}

				}
				return message;
			}
		});
	}

	@Override
	public void registerStompEndpoints(
			StompEndpointRegistry stompEndpointRegistry) {

		
		stompEndpointRegistry.addEndpoint("/ws").setAllowedOrigins("*")
				.withSockJS().setInterceptors(new HttpSessionHandshakeInterceptor());
	}

	@Override
	public void configureMessageBroker(MessageBrokerRegistry registry) {

		registry.enableSimpleBroker("/topic");// 3
	}

}
